import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # Ensure 3D plotting is registered

# 定义目标类
class Target:
    def __init__(self, initial_position, velocity):
        self.initial_position = np.array(initial_position)
        self.velocity = np.array(velocity)

    def get_position(self, time):
        """
        计算目标在给定时间的运动位置。
        :param time: 时间 t
        :return: 目标位置 (x, y, z)
        """
        return self.initial_position + self.velocity * time


# 定义雷达类
class Radar:
    def __init__(self, position, range):
        self.position = np.array(position)
        self.range = range

    def track(self, target_positions):
        """
        模拟雷达跟踪目标。
        :param target_positions: 目标位置列表 [(x, y, z), ...]
        :return: 雷达跟踪到的目标位置列表
        """
        tracked_positions = []
        for pos in target_positions:
            distance = np.linalg.norm(np.array(pos) - self.position)
            if distance <= self.range:
                tracked_positions.append(pos)
            else:
                tracked_positions.append(None)  # 超出雷达范围
        return tracked_positions


# 定义场景类
class SimulationScene:
    def __init__(self, radar, targets, simulation_time, time_step):
        self.radar = radar
        self.targets = targets
        self.simulation_time = simulation_time
        self.time_step = time_step

    def run(self):
        """
        运行仿真，生成目标轨迹和雷达跟踪结果。
        :return: 目标轨迹和雷达跟踪结果
        """
        time_points = np.arange(0, self.simulation_time + self.time_step, self.time_step)
        all_target_positions = [
            [target.get_position(t) for t in time_points] for target in self.targets
        ]
        radar_tracking_results = [
            self.radar.track(target_positions) for target_positions in all_target_positions
        ]
        return time_points, all_target_positions, radar_tracking_results


# 定义可视化类
class Visualization:
    @staticmethod
    def plot_3d(time_points, all_target_positions, radar_tracking_results):
        """
        绘制目标轨迹和雷达跟踪结果。
        """
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')

        # 绘制目标轨迹
        for i, target_positions in enumerate(all_target_positions):
            target_positions_array = np.array(target_positions)
            ax.plot(
                target_positions_array[:, 0],
                target_positions_array[:, 1],
                target_positions_array[:, 2],
                label=f"目标轨迹 {i + 1}",
                color='blue'
            )

        # 绘制雷达跟踪结果
        for i, tracking_results in enumerate(radar_tracking_results):
            tracked_positions_array = np.array(
                [pos if pos is not None else (np.nan, np.nan, np.nan) for pos in tracking_results]
            )
            ax.scatter(
                tracked_positions_array[:, 0],
                tracked_positions_array[:, 1],
                tracked_positions_array[:, 2],
                label=f"雷达跟踪 {i + 1}",
                color='red'
            )

        # 设置图形属性
        ax.set_xlabel("X (m)")
        ax.set_ylabel("Y (m)")
        # ax.set_zlabel("Z (m)")
        ax.set_title("三坐标雷达跟踪仿真")
        ax.legend()
        plt.show()


# 仿真参数
radar = Radar(position=(0, 0, 0), range=500)
targets = [
    Target(initial_position=(0, 0, 0), velocity=(50, 30, 10)),
    Target(initial_position=(100, 50, 20), velocity=(40, 20, 5))
]
simulation_time = 20  # 仿真时间 (s)
time_step = 1         # 时间步长 (s)

# 创建场景并运行仿真
scene = SimulationScene(radar=radar, targets=targets, simulation_time=simulation_time, time_step=time_step)
time_points, all_target_positions, radar_tracking_results = scene.run()

# 可视化结果
Visualization.plot_3d(time_points, all_target_positions, radar_tracking_results)